#coding:utf-8
import abc
from abc import ABC, ABCMeta, abstractmethod
from elasticsearch import Elasticsearch,helpers
import traceback
import numpy as np
import redis
import time
import requests
import json

if __name__ == "__main__":
    from setting import VSEARCH_HOST,VSEARCH_DB_FIELDS,LOG_PATH
    from log_helper import initMpTimedLogger
else:
    from .setting import VSEARCH_HOST,VSEARCH_DB_FIELDS,LOG_PATH
    from .log_helper import initMpTimedLogger

class RecallInterfaces(metaclass=ABCMeta):
    log = initMpTimedLogger(LOG_PATH, 'SimDocsRetrieval')

    def __init__(self, index):
        self.index = index
        self.ip_1 = '10.13.1.154'
        self.ip_2 = '10.13.1.160'
        self.ip_3 = '10.13.1.199'
        #self.ip_4 = '10.13.1.204'
        self.port = '9200'
        self.server = [
                {'host':self.ip_1,'port':self.port},
                {'host':self.ip_2,'port':self.port},
                {'host':self.ip_3,'port':self.port}
                #{'host':self.ip_4,'port':self.port}
                ]
        #elasticsearch 对象
        self.es = Elasticsearch(hosts=self.server, timeout=3)
    @abstractmethod
    def do_search(self, data, **kwargs):
        pass

    @abstractmethod
    def do_insert(self, data, **kwargs):
        pass
    
    @abstractmethod
    def do_delete(self, data, **kwargs):
        pass

    def term_search(self, value, size=20, fields='dataid'):
        ret = []
        if isinstance(value, list):
            body = {"query":{"terms":{fields:value}}}
        else:
            body = {"query":{"term":{fields:value}},"size":size}
        ret = self.es.search(self.index, body=body)['hits']['hits']
        return ret
    
    def es_search(self, body):
        '''
        range_search : {'size':500,'query':{'range':{'ctime':{'gte':begin,'lte':end}}}}
        '''
        result = self.es.search(self.index, body=body)#, preference='primary_first')
        return result['hits']['hits']

    def build(self, id, data, opType=''):
        '''build bulk data'''
        bdata = {'_index':self.index, '_id':id}
        if opType=='update':
            bdata['doc'] = data
            bdata['_op_type'] = 'update'
        elif opType=='delete':
            bdata['_op_type'] = 'delete'
        else:
            bdata['_source'] = data
        return bdata

    def bulk(self, bulkData):
        '''do bulk opeation'''
        try:
            return helpers.bulk(self.es, bulkData,request_timeout=3)
        except:
            return traceback.format_exc()

class RecallSimDocs():
    def __init__(self, method_object):
        self.method_object = method_object

    def do_search(self,data, **kwargs):
        return self.method_object.do_search(data, **kwargs)
    
    def do_insert(self,data, **kwargs):
        return self.method_object.do_insert(data, **kwargs)

    def do_delete(self,data, **kwargs):
        return self.method_object.do_delete(data, **kwargs)
    

class TextSearchClient(RecallInterfaces):
    def __init__(self, index):
        super(TextSearchClient, self).__init__(index)

    def do_search(self, data):
        assert "field" in data and 'title' in data
        count = data.get('count', 50)
        field = data['field']  # title or title_split
        q = data['title']

        body  = {
            "size":count,
            "query": {
                "bool": {
                "must": [{"match": {field: {"query":q,"operator": "OR","minimum_should_match": "20%"}}}],
                    #{"term": {"root": {"value": "true"}}}],
                "adjust_pure_negative":True,
                "boost": 1.0,
                }
            },
            "track_total_hits": 2147483647
        }
        time_begin = data.get('time_begin','')

        if time_begin and isinstance(time_begin, int):
            body["query"]["bool"]["filter"] = {"range":{"ctime":{"gte":time_begin}}}

        result = self.es_search(body)

        return result
    
    def do_delete(self,data):
        pass

    def do_insert(self, data):
        '''insert one data to es index'''
        n_bulkData = [self.build(data.get('dataid'), data, 'insert')]
        try:
            self.bulk(n_bulkData)
            return 1
        except Exception as err:
            return err
            

class VectorSearchClient(RecallInterfaces):
    '''
    功能：
        向量召回接口，支持search/insert/delete
        向量库milvus借助redis记录 milvus_id和group_id的映射关系。
    @index es索引名称
    @is_create_new_table 是否新建一个milvus collection, DEFAULT=False。 如果为True会先删除milvus库中该名字的collection，再创建。
    '''
    def __init__(self, index='group_id',host=VSEARCH_HOST):
        super(VectorSearchClient, self).__init__(index)
        self.vsearch_host = host     
        self.headers = {'Content-Type': 'application/json','Connection': 'close'}
        # TODO 换生产环境redis  向量入库存在延时，速度太快会产生召回失败问题。
        self.redis_client = redis.StrictRedis(host='10.23.5.198', port=6379, db=1, password='1qaz2wsx')

    def __post_request(self, data, opt):
        query_data = {}
        if opt == "search":
            query_data ={
                "data":{"embedding":data["embedding"]},
                "opt":"search"
            }
        elif opt == "insert":
            query_data ={
                "data":{},
                "opt":"insert"
            }
            for field in VSEARCH_DB_FIELDS:
                if field not in data:	continue
                if field == "group_id":
                    query_data["data"]["group_id"] = data[field]
                    continue
                query_data["data"][field] = data[field]
        elif opt == "delete":
            query_data ={
                "data":{"group_id":data["group_id"]},
                "opt":"delete"
            }
        else:
            return []
        try:
            response = requests.post(url=self.vsearch_host, 
                            headers=self.headers,
                            data=json.dumps(query_data)
                            )
            result = response.json()
            return result['ret']
        except Exception as e:
            print(e,response,query_data,result)
        return []

    def do_search(self, data, **kwargs):
        '''
        @data 查询向量 只支持单个向量搜索
        '''
        if not data or not isinstance(data, dict):
            return []
        query_vector = data.get("query_vector",[])
        if not query_vector or not isinstance(query_vector, list):
            return []
        return self.__post_request(data={"embedding":query_vector}, opt="search")

    def do_delete(self, data, **kwargs):
        if not data or not isinstance(data, dict):    return False
        if "dataid" not in data:    return False
        return self.__post_request(data, opt="delete")


    def do_insert(self, data, **kwargs):
        if not data or not isinstance(data, dict):
            return False
        return self.__post_request(data, opt="insert")

    def get_vec(self, ids):
        pass

    def get_newsid_by_ids(self, ids):
        # 用milvus id 查 news_id 
        pass

    def get_ids_by_newsid(self, news_ids):
        # 用news_id 查 milvus
        pass

    def get_docs(self, ids, id_type="news_id"):
        pass



if __name__ == "__main__":
    #ids = [1622537878397089024]
    vsc = VectorSearchClient()
    esc = TextSearchClient('group_id_v0')
    RSClient = RecallSimDocs(esc)

    data = {"field":"title_split", "title":"表情 大师 阿扎尔 代言 比利时 麦当劳 汉堡"}
    data = {"field":"title_split", "title":"本色 出演 阿扎尔 牛肉 汉堡 做广告 引发 热议"}
    #data = {"field":"title_ori", "title":"本色出演阿扎尔牛肉汉堡做广告引发热议"}

    _st1 = time.time()
    ret = RSClient.do_search(data)
    _st2 = time.time()
    for r in ret:
        print(r['_source']['title_ori'].strip(), r['_source']['title_fp'])

    
    vec = [-0.2395342, 0.2476821, -0.0182051, -0.1149711, 0.2403101, -0.11299, 0.0893538, 0.0773847, -0.0838335, 0.1837936, 0.1655515, 0.3418716, 0.1083572, -0.0357101, -0.0037128, 0.2391805, -0.0538993, 0.0073957, -0.081176, 0.0099233, 0.1095814, 0.0392749, -0.0647294, 0.1923988, -0.252735, 0.3340638, 0.2728182, -0.0049818, 0.0407264, 0.3632605, 0.0421063, 0.2796178]
    vec = [0.13459999859333038, 0.03759999945759773, 0.11959999799728394, -0.029500000178813934, -0.17100000381469727, 0.0024999999441206455, -0.03750000149011612, 0.01140000019222498, -0.09160000085830688, -0.20100000500679016, 0.08209999650716782, -0.17030000686645508, -0.015300000086426735, -0.08649999648332596, 0.27559998631477356, 0.04410000145435333, 0.08669999986886978, 0.0012000000569969416, 0.15809999406337738, 0.22310000658035278, 0.2004999965429306, -0.1881999969482422, 0.012799999676644802, -0.09120000153779984, -0.08190000057220459, -0.0006000000284984708, -0.06210000067949295, 0.0835999995470047, -0.09960000216960907, 0.17139999568462372, 0.0034000000450760126, -0.11810000240802765, -0.3224000036716461, -0.04610000178217888, -0.1088000014424324, 0.0003000000142492354, 0.10679999738931656, 0.02930000051856041, -0.19580000638961792, -0.10890000313520432, -0.1395999938249588, 0.022099999710917473, 0.03579999879002571, 0.009800000116229057, 0.1242000013589859, 0.06859999895095825, -0.04729999974370003, 0.14560000598430634, 0.05570000037550926, 0.12839999794960022, 0.03460000082850456, -0.04520000144839287, 0.003800000064074993, 0.09359999746084213, -0.027300000190734863, 0.13770000636577606, 0.24459999799728394, 0.11649999767541885, 0.1501999944448471, 0.22100000083446503, 0.06859999895095825, -0.24740000069141388, 0.12389999628067017, -0.066600002348423]
    vec = [0.10450000315904617, -0.10559999942779541, 0.019600000232458115, 0.023000000044703484, -0.047200001776218414, 0.10999999940395355, -0.027699999511241913, 0.05400000140070915, -0.274399995803833, -0.2939999997615814, 0.12470000237226486, -0.007600000128149986, -0.004399999976158142, 0.23469999432563782, -0.0723000019788742, -0.15760000050067902, 0.15889999270439148, 0.018200000748038292, 0.26570001244544983, 0.011900000274181366, 0.10920000076293945, -0.044199999421834946, -0.13050000369548798, -0.12479999661445618, -0.06650000065565109, 0.023000000044703484, -0.12479999661445618, -0.01730000041425228, -0.20469999313354492, 0.16339999437332153, 0.05220000073313713, -0.18649999797344208, -0.23409999907016754, 0.06520000100135803, -0.0414000004529953, 0.06729999929666519, 0.1046999990940094, -0.15770000219345093, -0.057100001722574234, 0.03539999946951866, 0.08449999988079071, -0.06379999965429306, -0.05770000070333481, -0.03999999910593033, -0.09610000252723694, 0.2289000004529953, -0.08940000087022781, 0.0012000000569969416, 0.12039999663829803, -0.09619999676942825, 0.20180000364780426, 0.029200000688433647, -0.10679999738931656, -0.13449999690055847, 0.04749999940395355, -0.1526000052690506, 0.15950000286102295, -0.13120000064373016, 0.14640000462532043, 0.07559999823570251, -0.09539999812841415, -0.05249999836087227, -0.07720000296831131, 0.16539999842643738]
    vec = [-0.051500000059604645, -0.020099999383091927, 0.06199999898672104, 0.007600000128149986, 0.20839999616146088, 0.040699999779462814, -0.13079999387264252, -0.12720000743865967, -0.07500000298023224, -0.2224999964237213, 0.0575999990105629, -0.05290000140666962, 0.05779999867081642, 0.09700000286102295, 0.035999998450279236, -0.26269999146461487, -0.17949999868869781, -0.07810000330209732, 0.04830000177025795, -0.12449999898672104, -0.06669999659061432, -0.025100000202655792, 0.1339000016450882, -0.08560000360012054, -0.12129999697208405, 0.009499999694526196, -0.0478999987244606, -0.03680000081658363, 0.014600000344216824, 0.09650000184774399, 0.025699999183416367, -0.04910000041127205, -0.11670000106096268, 0.03689999878406525, -0.10580000281333923, -0.2768999934196472, -0.09719999879598618, -0.07349999994039536, -0.15880000591278076, -0.10830000042915344, 0.021299999207258224, 0.28949999809265137, -0.09960000216960907, 0.05350000038743019, -0.07460000365972519, 0.04439999908208847, -0.13699999451637268, -0.050700001418590546, 0.2581999897956848, 0.12330000102519989, -0.21789999306201935, 0.18469999730587006, 0.1768999993801117, 0.10199999809265137, 0.06069999933242798, -0.11420000344514847, 0.02590000070631504, 0.08460000157356262, -0.0803999975323677, -0.18320000171661377, -0.15039999783039093, -0.05389999970793724, 0.22949999570846558, -0.15600000321865082]
    RSClient.method_object = vsc

    _st3 = time.time()
    data = {"query_vector":vec}
    ret = RSClient.do_search(data)
    _st4 = time.time()

    for r in ret:
        #print(r['_source']['title_ori'].strip(), r['_source']['title_fp'])
        print(r)

    print(1000*(_st2-_st1), 1000*(_st3-_st2), 1000*(_st4-_st3))

